package com.gitee.wsl.flow.coroutines

import com.gitee.wsl.coroutines.strategy.lazy.LazyDeferred
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineDispatcher
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.ExperimentalCoroutinesApi
import kotlinx.coroutines.Job
import kotlinx.coroutines.cancel
import kotlinx.coroutines.channels.ReceiveChannel
import kotlinx.coroutines.channels.produce
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.FlowCollector
import kotlinx.coroutines.flow.flatMapMerge
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.flow.flowOn
import kotlinx.coroutines.flow.map

class AsyncFlow<out T>(
    val deferredFlow: Flow<LazyDeferred<T>>,
) : Flow<T> {
    override suspend fun collect(collector: FlowCollector<T>): Unit =
        deferredFlow.collect { collector.emit((it.await())) }
}

inline fun <T, R> Flow<T>.async(
    dispatcher: CoroutineDispatcher = Dispatchers.Default,
    crossinline block: suspend CoroutineScope.(T) -> R,
): AsyncFlow<R> {
    val flow = map { LazyDeferred(dispatcher) { block(it) } }
    return AsyncFlow(flow)
}

inline fun <T, R> AsyncFlow<T>.map(crossinline action: (T) -> R): AsyncFlow<R> =
    AsyncFlow(deferredFlow.map { input ->
        //TODO add function composition
        LazyDeferred(input.dispatcher) {
            input.start(this)
            action(input.await())
        }
    })

@OptIn(ExperimentalCoroutinesApi::class)
suspend fun <T> AsyncFlow<T>.collect(concurrency: Int, collector: FlowCollector<T>) {
    require(concurrency >= 1) { "Buffer size should be more than 1, but was $concurrency" }

    coroutineScope {
        //Starting up to N deferred coroutines ahead of time
        val channel: ReceiveChannel<LazyDeferred<T>> = produce(capacity = concurrency - 1) {
            deferredFlow.collect { value ->
                value.start(this@coroutineScope)
                send(value)
            }
        }

        (channel as Job).invokeOnCompletion {
            if (it is CancellationException && it.cause == null) cancel()
        }

        for (element in channel) {
            collector.emit(element.await())
        }

        val producer = channel as Job
        if (producer.isCancelled) {
            producer.join()
            //throw producer.getCancellationException()
        }
    }
}

suspend inline fun <T> AsyncFlow<T>.collect(
    concurrency: Int,
    crossinline action: suspend (value: T) -> Unit,
): Unit = collect(concurrency, FlowCollector<T> { value -> action(value) })

@OptIn(ExperimentalCoroutinesApi::class)
inline fun <T, R> Flow<T>.mapParallel(
    dispatcher: CoroutineDispatcher = Dispatchers.Default,
    crossinline transform: suspend (T) -> R,
): Flow<R> = flatMapMerge { value -> flow { emit(transform(value)) } }.flowOn(dispatcher)
